package br.com.linkin.mi.utils.xml;

import br.com.linkin.mi.utils.IO;
import br.com.linkin.mi.utils.security.Security;
import java.io.*;
import java.net.URL;
import java.security.*;
import java.security.cert.X509Certificate;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import javax.xml.XMLConstants;
import javax.xml.bind.JAXBContext;
import javax.xml.bind.JAXBElement;
import javax.xml.bind.Marshaller;
import javax.xml.bind.Unmarshaller;
import javax.xml.crypto.MarshalException;
import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.SAXParserFactory;
import javax.xml.transform.*;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import javax.xml.transform.stream.StreamSource;
import javax.xml.validation.Schema;
import javax.xml.validation.SchemaFactory;
import javax.xml.validation.Validator;
import javax.xml.xpath.XPathFactory;
import org.w3c.dom.*;
import org.w3c.dom.ls.LSInput;
import org.xml.sax.InputSource;

public class Xml {

    /**
     * OutputJaxpImplementationInfo
     */
    public static void OutputJaxpImplementationInfo() {
        System.out.println(getJaxpImplementationInfo("DocumentBuilderFactory",
                DocumentBuilderFactory.newInstance().getClass()));
        System.out.println(getJaxpImplementationInfo("XPathFactory",
                XPathFactory.newInstance().getClass()));
        System.out.println(getJaxpImplementationInfo("TransformerFactory",
                TransformerFactory.newInstance().getClass()));
        System.out.println(getJaxpImplementationInfo("SAXParserFactory",
                SAXParserFactory.newInstance().getClass()));
    }

    /**
     * getJaxpImplementationInfo
     * 
     * @param componentName
     * @param componentClass
     * @return 
     */
    public static String getJaxpImplementationInfo(String componentName,
            Class componentClass) {
        CodeSource source = componentClass.getProtectionDomain().getCodeSource();
        return MessageFormat.format("{0} implementation: {1} loaded from: {2}",
                componentName, componentClass.getName(),
                source == null ? "Java Runtime" : source.getLocation());
    }

    /**
     * getFirstTagValue
     * 
     * @param doc
     * @param tagName
     * @return 
     */
    public static String getFirstTagValue(Document doc, String tagName) {
        String retorno = null;
        Node node = getFirstElementByTag(doc, tagName);
        if (node != null) {
            retorno = node.getTextContent();
        }
        return retorno;
    }

    /**
     * getFirstElementByTag
     * 
     * @param doc
     * @param tagName
     * @return 
     */
    public static Node getFirstElementByTag(Document doc, String tagName) {
        Node retorno = null;
        NodeList nodeList = doc.getElementsByTagName(tagName);
        if (nodeList.getLength() > 0) {
            retorno = nodeList.item(0);
        }
        return retorno;
    }

    /**
     * createJAXBContext
     * 
     * @param packages
     * @return
     * @throws Exception 
     */
    public static JAXBContext createJAXBContext(String... packages)
            throws Exception {
        String pack = "";
        for (String packName : packages) {
            pack = pack + ("".equals(pack) ? "" : ":") + packName;
        }
        return JAXBContext.newInstance(pack);
    }

    /**
     * createMarshaller
     * 
     * @param packages
     * @return
     * @throws Exception 
     */
    public static Marshaller createMarshaller(String... packages)
            throws Exception {
        return createJAXBContext(packages).createMarshaller();
    }

    /**
     * createUnmarshaller
     * 
     * @param packages
     * @return
     * @throws Exception 
     */
    public static Unmarshaller createUnmarshaller(String... packages)
            throws Exception {
        return createJAXBContext(packages).createUnmarshaller();
    }

    /**
     * createMarshaller
     * 
     * @param packageClass
     * @return
     * @throws Exception 
     */
    private static Marshaller createMarshaller(Class<?> packageClass)
            throws Exception {
        String pack = packageClass.getPackage().getName();
        JAXBContext jc1 = JAXBContext.newInstance(pack);
        return jc1.createMarshaller();
    }

    /**
     * createUnmarshaller
     * 
     * @param packageClass
     * @return
     * @throws Exception 
     */
    public static Unmarshaller createUnmarshaller(Class<?>... packageClass)
            throws Exception {
        String pack = "";
        for (Class<?> clazz : packageClass) {
            pack = pack + ("".equals(pack) ? "" : ":")
                    + clazz.getPackage().getName();
        }
        JAXBContext jc1 = JAXBContext.newInstance(pack);
        return jc1.createUnmarshaller();
    }

    /**
     * marshalToString
     * 
     * @param obj
     * @return
     * @throws Exception 
     */
    public static String marshalToString(Object obj) throws Exception {
        StringWriter str = new StringWriter();
        marshal(str, obj);
        return str.toString();
    }

    /**
     * marshallToDocument
     * 
     * @param obj
     * @return
     * @throws Exception 
     */
    public static Document marshallToDocument(Object obj) throws Exception {
        return marshallToDocument(obj, true);
    }

    /**
     * marshallToDocument
     * 
     * @param obj
     * @param nameSpaceAwarea
     * @return
     * @throws Exception 
     */
    public static Document marshallToDocument(Object obj,
            boolean nameSpaceAwarea) throws Exception {
        Marshaller marshaller = createMarshaller(obj.getClass());
        Document doc = Xml.createDocument(nameSpaceAwarea);
        marshaller.marshal(obj, doc);
        return doc;
    }

    /**
     * marshallToDocumentAndFixNamespace
     * 
     * @param obj
     * @return
     * @throws Exception 
     */
    public static Document marshallToDocumentAndFixNamespace(Object obj)
            throws Exception {
        return marshallToDocumentAndFixNamespace(obj, true);
    }

    /**
     * marshallToDocumentAndFixNamespace
     * 
     * @param obj
     * @param nameSpaceAware
     * @return
     * @throws Exception 
     */
    public static Document marshallToDocumentAndFixNamespace(Object obj,
            boolean nameSpaceAware) throws Exception {
        Document doc = marshallToDocument(obj, nameSpaceAware);
        removeNS2Namespace(doc.getDocumentElement());
        return doc;
    }

    /**
     * marshallToDocument
     * 
     * @param marshaller
     * @param obj
     * @return
     * @throws Exception 
     */
    public static Document marshallToDocument(Marshaller marshaller, Object obj)
            throws Exception {
        Document doc = Xml.createDocument();
        marshaller.marshal(obj, doc);
        return doc;
    }

    /**
     * printTo
     * 
     * @param doc
     * @param out
     * @param omitXmlDeclaration
     * @throws Exception 
     */
    public static void printTo(Document doc, Writer out,
            boolean omitXmlDeclaration) throws Exception {
        DOMSource domSrc = new DOMSource(doc);
        StreamResult streamResult = new StreamResult(out);
        TransformerFactory transformerFac = createTransformerFactory();
        Transformer transformer = transformerFac.newTransformer();
        if (omitXmlDeclaration) {
            transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION,
                    "yes");
        }
        transformer.transform(domSrc, streamResult);
    }

    /**
     * createTransformerFactory
     * 
     * @return
     * @throws TransformerFactoryConfigurationError 
     */
    public static TransformerFactory createTransformerFactory()
            throws TransformerFactoryConfigurationError {
        return TransformerFactory.newInstance(
                "org.apache.xalan.processor.TransformerFactoryImpl", Thread.currentThread().getContextClassLoader());
    }

    /**
     * printToString
     * 
     * @param doc
     * @return
     * @throws Exception 
     */
    public static String printToString(Document doc) throws Exception {
        StringWriter str = new StringWriter();
        printTo(doc, str, false);
        return str.toString();
    }

    /**
     * printTo
     * 
     * @param doc
     * @param out
     * @throws Exception 
     */
    public static void printTo(Document doc, OutputStream out) throws Exception {
        DOMSource domSrc = new DOMSource(doc);
        StreamResult streamResult = new StreamResult(out);
        TransformerFactory transformerFac = createTransformerFactory();
        Transformer transformer = transformerFac.newTransformer();
        transformer.transform(domSrc, streamResult);
    }

    /**
     * createDocument
     * 
     * @return
     * @throws Exception 
     */
    public static Document createDocument() throws Exception {
        return createDocument(true);
    }

    /**
     * createDocument
     * 
     * @param namespaceAware
     * @return
     * @throws Exception 
     */
    public static Document createDocument(boolean namespaceAware)
            throws Exception {
        DocumentBuilderFactory dbFactory = createDocumentBuilderFactory();
        dbFactory.setNamespaceAware(namespaceAware);
        DocumentBuilder builder = dbFactory.newDocumentBuilder();
        return builder.newDocument();
    }

    /**
     * createDocumentBuilderFactory
     * 
     * @return 
     */
    public static DocumentBuilderFactory createDocumentBuilderFactory() {
        DocumentBuilderFactory dbFactory = DocumentBuilderFactory.newInstance();
        return dbFactory;
    }

    /**
     * sign
     * 
     * @param keyStore
     * @param password
     * @param doc
     * @param referenceUri
     * @throws Exception 
     */
    public static void sign(String keyStore, String password, Document doc,
            String referenceUri) throws Exception {
        sign(keyStore, password, doc.getDocumentElement(), referenceUri);
    }

    /**
     * sign
     * 
     * @param keyStore
     * @param password
     * @param node
     * @param referenceUri
     * @throws Exception 
     */
    public static void sign(String keyStore, String password, Node node,
            String referenceUri) throws Exception {
        Document retorno = null;
        KeyStore.PrivateKeyEntry key = Security.loadPkcs12PrivateKey(
                keyStore, password);
        sign(key, node, referenceUri);
    }

    /**
     * sign
     * 
     * @param key
     * @param node
     * @param referenceUri
     * @throws NoSuchAlgorithmException
     * @throws InvalidAlgorithmParameterException
     * @throws MarshalException
     * @throws XMLSignatureException 
     */
    private static void sign(KeyStore.PrivateKeyEntry key, Node node,
            String referenceUri) throws NoSuchAlgorithmException,
            InvalidAlgorithmParameterException, MarshalException,
            XMLSignatureException {
        DOMSignContext dsc = new DOMSignContext(key.getPrivateKey(), node);
        XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");
        List transforms = new ArrayList();
        transforms.add(fac.newTransform(Transform.ENVELOPED,
                (TransformParameterSpec) null));
        transforms.add(fac.newTransform(
                "http://www.w3.org/TR/2001/REC-xml-c14n-20010315",
                (TransformParameterSpec) null));
        Reference ref = fac.newReference('#' + referenceUri,
                fac.newDigestMethod(DigestMethod.SHA1, null), transforms, null,
                null);

        SignedInfo si = fac.newSignedInfo(fac.newCanonicalizationMethod(
                CanonicalizationMethod.INCLUSIVE,
                (C14NMethodParameterSpec) null), fac.newSignatureMethod(SignatureMethod.RSA_SHA1, null),
                Collections.singletonList(ref));

        X509Certificate cert = (X509Certificate) key.getCertificate();

        KeyInfoFactory kif = fac.getKeyInfoFactory();
        List x509Content = new ArrayList();
        x509Content.add(cert);
        X509Data xd = kif.newX509Data(x509Content);
        KeyInfo ki = kif.newKeyInfo(Collections.singletonList(xd));
        XMLSignature signature = fac.newXMLSignature(si, ki);
        signature.sign(dsc);

        // DOMResult signedNfe = new DOMResult();
        // TransformerFactory tf = TransformerFactory.newInstance();
        // Transformer trans = tf.newTransformer();
        // trans.transform(new DOMSource(doc), signedNfe);
        // retorno = (Document) signedNfe.getNode();
    }

    /**
     * marshal
     * 
     * @param out
     * @param obj
     * @throws Exception 
     */
    public static void marshal(Writer out, Object obj) throws Exception {
        Marshaller marshaller = createMarshaller(obj.getClass());
        marshaller.marshal(obj, out);
    }

    /**
     * marshal
     * 
     * @param doc
     * @param obj
     * @throws Exception 
     */
    public static void marshal(Document doc, Object obj) throws Exception {
        Marshaller marshaller = createMarshaller(obj.getClass());
        marshaller.marshal(obj, doc);
    }

    /**
     * unmarshal
     * 
     * @param in
     * @param clazz
     * @return
     * @throws Exception 
     */
    public static Object unmarshal(Node in, Class<?>... clazz) throws Exception {
        Unmarshaller unmarshaller = createUnmarshaller(clazz);
        Object retorno = unmarshaller.unmarshal(in);
        if (retorno instanceof JAXBElement) {
            retorno = ((JAXBElement<?>) retorno).getValue();
        }
        return retorno;
    }

    /**
     * unmarshal
     * 
     * @param in
     * @param clazz
     * @return
     * @throws Exception 
     */
    public static Object unmarshal(Document in, Class<?>... clazz)
            throws Exception {
        Unmarshaller unmarshaller = createUnmarshaller(clazz);
        Object retorno = unmarshaller.unmarshal(in);
        if (retorno instanceof JAXBElement) {
            retorno = ((JAXBElement<?>) retorno).getValue();
        }
        return retorno;
    }

    /**
     * unmarshal
     * 
     * @param in
     * @param clazz
     * @return
     * @throws Exception 
     */
    public static Object unmarshal(Reader in, Class<?>... clazz)
            throws Exception {
        Unmarshaller unmarshaller = createUnmarshaller(clazz);
        Object retorno = unmarshaller.unmarshal(in);
        if (retorno instanceof JAXBElement) {
            retorno = ((JAXBElement<?>) retorno).getValue();
        }
        return retorno;
    }

    /**
     * sign
     * 
     * @param obj
     * @param keyStore
     * @param keyStorePassword
     * @param referenceUri
     * @throws Exception 
     */
    public static void sign(Document obj, InputStream keyStore,
            String keyStorePassword, String referenceUri) throws Exception {
        KeyStore.PrivateKeyEntry key = Security.loadPrivateKey(
                Security.loadPkcsKeystore(keyStore, keyStorePassword),
                keyStorePassword);
        sign(key, obj.getDocumentElement(), referenceUri);

    }

    /**
     * removeNS2Namespace
     * 
     * @param element 
     */
    public static void removeNS2Namespace(Element element) {
        element.removeAttribute("xmlns:ns2");
        element.removeAttribute("xmlns:ns3");
    }

    /**
     * addOrSetAttribute
     * 
     * @param doc
     * @param tagName
     * @param attName
     * @param attValue 
     */
    public static void addOrSetAttribute(Document doc, String tagName,
            String attName, String attValue) {
        NodeList nl = doc.getElementsByTagName(tagName);
        Element node = null;
        for (int i = 0; i < nl.getLength(); i++) {
            node = (Element) nl.item(i);
            node.setAttribute(attName, attValue);
        }
    }

    /**
     * removeAllNamespaceAttributes
     * 
     * @param element 
     */
    public static void removeAllNamespaceAttributes(Element element) {
        NamedNodeMap nodeMap = element.getAttributes();
        List<Node> nodeList = new ArrayList<Node>();
        for (int i = 0; i < nodeMap.getLength(); i++) {
            Node att = nodeMap.item(i);
            if (att.getNodeName().contains("xmlns")) {
                nodeList.add(att);
            }
        }
        for (Node node : nodeList) {
            element.removeAttribute(node.getNodeName());
        }
    }

    /**
     * merge
     * 
     * @param target
     * @param origin
     * @return 
     */
    public static Document merge(Document target, Document origin) {
        return merge(target, origin, target.getDocumentElement());
    }

    /**
     * merge
     * 
     * @param target
     * @param origin
     * @param element
     * @return 
     */
    public static Document merge(Document target, Document origin, Node element) {
        Node node = target.importNode(origin.getDocumentElement(), true);
        element.appendChild(node);
        return target;
    }

    /**
     * loadDocument
     * 
     * @param str
     * @return
     * @throws Exception 
     */
    public static Document loadDocument(String str) throws Exception {
        return loadDocument(str, false);
    }

    /**
     * loadDocument
     * 
     * @param str
     * @param namespaceAware
     * @return
     * @throws Exception 
     */
    public static Document loadDocument(String str, boolean namespaceAware)
            throws Exception {
        DocumentBuilderFactory dbFactory = createDocumentBuilderFactory();
        dbFactory.setNamespaceAware(namespaceAware);
        dbFactory.setValidating(false);
        dbFactory.setIgnoringComments(true);
        DocumentBuilder builder = dbFactory.newDocumentBuilder();
        InputSource source = new InputSource(new StringReader(str));
        return builder.parse(source);
    }

    /**
     * loadDocument
     * 
     * @param str
     * @return
     * @throws Exception 
     */
    public static Document loadDocument(InputStream str) throws Exception {
        return loadDocument(str, true);
    }

    /**
     * loadDocument
     * 
     * @param str
     * @param namespaceAware
     * @return
     * @throws Exception 
     */
    public static Document loadDocument(InputStream str, boolean namespaceAware)
            throws Exception {
        DocumentBuilderFactory dbFactory = createDocumentBuilderFactory();
        dbFactory.setNamespaceAware(namespaceAware);
        dbFactory.setValidating(false);
        dbFactory.setIgnoringComments(true);
        DocumentBuilder builder = dbFactory.newDocumentBuilder();
        InputSource source = new InputSource(str);
        return builder.parse(source);
    }

    /**
     * isSignatureValid
     * 
     * @param key
     * @param doc
     * @return
     * @throws Exception 
     */
    public static boolean isSignatureValid(Key key, Document doc)
            throws Exception {
        XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");
        NodeList nl = doc.getElementsByTagNameNS(XMLSignature.XMLNS,
                "Signature");
        if (nl.getLength() == 0) {
            throw new Exception("Cannot find Signature element");
        }
        DOMValidateContext valContext = new DOMValidateContext(key, nl.item(0));

        XMLSignature signature = fac.unmarshalXMLSignature(valContext);
        return signature.validate(valContext);
    }

    /**
     * SourceIdentifier
     */
    public static class SourceIdentifier {

        String sourceURI;
        String sourceName;
        Source source;

        public String getSourceURI() {
            return sourceURI;
        }

        public void setSourceURI(String sourceURI) {
            this.sourceURI = sourceURI;
        }

        public String getSourceName() {
            return sourceName;
        }

        public void setSourceName(String sourceName) {
            this.sourceName = sourceName;
        }

        public Source getSource() {
            return source;
        }

        public void setSource(Source source) {
            this.source = source;
        }
    }

    /**
     * validate
     * 
     * @param doc
     * @param sources
     * @throws Exception 
     */
    public static void validate(Document doc, final Xml.SourceIdentifier[] sources)
            throws Exception {
        SchemaFactory factory = SchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
        // factory.setResourceResolver(new LSResourceResolver() {
        // public LSInput resolveResource(String type, String namespaceURI,
        // String publicId, String systemId, String baseURI) {
        // InputStream in = null;
        // for (SourceIdentifier si : sources) {
        // if (systemId.equals(si.getSourceName())) {
        // in = IO.load(si.getSourceURI());
        // break;
        // }
        // }
        // return new LSInputImpl(publicId, systemId, in);
        // }
        // });
        Source[] sourceList = new Source[sources.length];
        int i = 0;
        for (Xml.SourceIdentifier sourceIdentifier : sources) {
            sourceList[i] = sourceIdentifier.getSource();
            i++;
        }
        Schema schema = factory.newSchema(sourceList);
        Validator validator = schema.newValidator();
        validator.validate(new DOMSource(doc));
    }

    /**
     * newValidator
     * 
     * @param resourceName
     * @return
     * @throws Exception 
     */
    public static Validator newValidator(String resourceName) throws Exception {
        SchemaFactory factory = SchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
        URL resourceURL = IO.loadURL(resourceName);
        StreamSource source = new StreamSource(IO.load(resourceName),
                resourceURL.toURI().toString());
        Schema schema = factory.newSchema(source);
        Validator validator = schema.newValidator();
        return validator;

    }

    /**
     * validate
     * 
     * @param doc
     * @param resourceName
     * @throws Exception 
     */
    public static void validate(Node doc, String resourceName) throws Exception {
        Validator validator = newValidator(resourceName);
        validator.validate(new DOMSource(doc));
    }

    /**
     * validate
     * 
     * @param doc
     * @param source
     * @throws Exception 
     */
    public static void validate(Document doc, final Source source)
            throws Exception {
        SchemaFactory factory = SchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
        // factory.setResourceResolver(new LSResourceResolver() {
        // public LSInput resolveResource(String type, String namespaceURI,
        // String publicId, String systemId, String baseURI) {
        // URI uri = new URI();
        // return new LSInputImpl(publicId, systemId,
        // IO.load(source.getSystemId()));
        // }
        // });

        Schema schema = factory.newSchema(source);
        Validator validator = schema.newValidator();
        validator.validate(new DOMSource(doc));
    }

    /**
     * LSInputImpl
     */
    static class LSInputImpl implements LSInput {

        private String publicId;
        private String systemId;

        public String getPublicId() {
            return publicId;
        }

        public void setPublicId(String publicId) {
            this.publicId = publicId;
        }

        public String getBaseURI() {
            return null;
        }

        public InputStream getByteStream() {
            return null;
        }

        public boolean getCertifiedText() {
            return false;
        }

        public Reader getCharacterStream() {
            return null;
        }

        public String getEncoding() {
            return null;
        }

        public String getStringData() {
            synchronized (inputStream) {
                try {
                    byte[] input = new byte[inputStream.available()];
                    inputStream.read(input);
                    String contents = new String(input);
                    return contents;
                } catch (IOException e) {
                    e.printStackTrace();
                    System.out.println("Exception " + e);
                    return null;
                }
            }
        }

        public void setBaseURI(String baseURI) {
        }

        public void setByteStream(InputStream byteStream) {
        }

        public void setCertifiedText(boolean certifiedText) {
        }

        public void setCharacterStream(Reader characterStream) {
        }

        public void setEncoding(String encoding) {
        }

        public void setStringData(String stringData) {
        }

        public String getSystemId() {
            return systemId;
        }

        public void setSystemId(String systemId) {
            this.systemId = systemId;
        }

        public BufferedInputStream getInputStream() {
            return inputStream;
        }

        public void setInputStream(BufferedInputStream inputStream) {
            this.inputStream = inputStream;
        }
        private BufferedInputStream inputStream;

        public LSInputImpl(String publicId, String sysId, InputStream input) {
            this.publicId = publicId;
            this.systemId = sysId;
            this.inputStream = new BufferedInputStream(input);
        }
    }
}
